"""
07_data_loader.py
=================
Healthcare Knowledge Graph Data Loader

This module covers:
1. Loading JSON data into Neo4j
2. Creating healthcare entities (Patient, Disease, Medication)
3. Establishing relationships
4. Bulk import patterns
5. Data validation

Author: Neo4j Tutorial
Date: January 2025
"""

import json
from typing import Dict, List, Any
from pathlib import Path
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class HealthcareDataLoader:
    """
    Loads healthcare data into Neo4j knowledge graph.
    
    Schema Overview:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │              HEALTHCARE KNOWLEDGE GRAPH SCHEMA                  │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │                        ┌──────────┐                            │
    │                        │ Disease  │                            │
    │                        │ (ICD-11) │                            │
    │                        └────▲─────┘                            │
    │                             │                                   │
    │                    DIAGNOSED_WITH                               │
    │                             │                                   │
    │  ┌──────────┐         ┌────┴─────┐         ┌──────────┐       │
    │  │Medication│◄─TAKES──│ Patient  │──ALLERGIC_TO──►│ Allergy  │ │
    │  │ (RxNorm) │         └──────────┘                └──────────┘ │
    │  └────┬─────┘                                                   │
    │       │                                                         │
    │    TREATS                                                       │
    │       │                                                         │
    │       ▼                                                         │
    │  ┌──────────┐                                                   │
    │  │ Disease  │                                                   │
    │  └──────────┘                                                   │
    │                                                                 │
    │  Relationship Properties:                                       │
    │  - DIAGNOSED_WITH: diagnosedDate, diagnosedBy, status          │
    │  - TAKES: dosage, frequency, prescribedDate, status            │
    │  - TREATS: effectiveness, firstLine                            │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, driver=None):
        self.driver = driver or get_driver()
    
    def clear_database(self):
        """Clear all data from the database."""
        query = """
        MATCH (n)
        DETACH DELETE n
        """
        with self.driver.session() as session:
            session.run(query)
        print("🗑️  Database cleared")
    
    def create_constraints(self):
        """Create unique constraints for all entity types."""
        constraints = [
            ("Patient", "id"),
            ("Disease", "id"),
            ("Medication", "id"),
            ("Document", "id"),
            ("Chunk", "id"),
        ]
        
        with self.driver.session() as session:
            for label, prop in constraints:
                query = f"""
                CREATE CONSTRAINT constraint_{label.lower()}_{prop}_unique 
                IF NOT EXISTS
                FOR (n:{label})
                REQUIRE n.{prop} IS UNIQUE
                """
                try:
                    session.run(query)
                    print(f"   ✓ Constraint on {label}.{prop}")
                except Exception as e:
                    print(f"   ⚠ Constraint on {label}.{prop}: {e}")
    
    def create_indexes(self):
        """Create indexes for common query patterns."""
        indexes = [
            ("Patient", "name"),
            ("Patient", "gender"),
            ("Disease", "category"),
            ("Medication", "drugClass"),
        ]
        
        with self.driver.session() as session:
            for label, prop in indexes:
                query = f"""
                CREATE INDEX idx_{label.lower()}_{prop.lower()} 
                IF NOT EXISTS
                FOR (n:{label})
                ON (n.{prop})
                """
                try:
                    session.run(query)
                    print(f"   ✓ Index on {label}.{prop}")
                except Exception as e:
                    print(f"   ⚠ Index on {label}.{prop}: {e}")
    
    # =========================================================
    # ENTITY CREATION
    # =========================================================
    
    def load_patients(self, patients: List[Dict]) -> int:
        """Load patient data."""
        query = """
        UNWIND $patients AS p
        MERGE (patient:Patient {id: p.id})
        SET patient.name = p.name,
            patient.dateOfBirth = date(p.dateOfBirth),
            patient.gender = p.gender,
            patient.bloodType = p.bloodType,
            patient.allergies = p.allergies,
            patient.active = p.active,
            patient.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, patients=patients)
            return result.single()["created"]
    
    def load_diseases(self, diseases: List[Dict]) -> int:
        """Load disease data (ICD-11 codes)."""
        query = """
        UNWIND $diseases AS d
        MERGE (disease:Disease {id: d.id})
        SET disease.name = d.name,
            disease.category = d.category,
            disease.description = d.description,
            disease.severity = d.severity,
            disease.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, diseases=diseases)
            return result.single()["created"]
    
    def load_medications(self, medications: List[Dict]) -> int:
        """Load medication data (RxNorm codes)."""
        query = """
        UNWIND $medications AS m
        MERGE (med:Medication {id: m.id})
        SET med.name = m.name,
            med.genericName = m.genericName,
            med.drugClass = m.drugClass,
            med.dosageForms = m.dosageForms,
            med.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, medications=medications)
            return result.single()["created"]
    
    def load_documents(self, documents: List[Dict]) -> int:
        """Load document data for RAG."""
        query = """
        UNWIND $documents AS d
        MERGE (doc:Document {id: d.id})
        SET doc.title = d.title,
            doc.source = d.source,
            doc.content = d.content,
            doc.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, documents=documents)
            return result.single()["created"]
    
    # =========================================================
    # RELATIONSHIP CREATION
    # =========================================================
    
    def load_diagnoses(self, diagnoses: List[Dict]) -> int:
        """Load diagnosis relationships."""
        query = """
        UNWIND $diagnoses AS d
        MATCH (p:Patient {id: d.patientId})
        MATCH (disease:Disease {id: d.diseaseId})
        MERGE (p)-[r:DIAGNOSED_WITH]->(disease)
        SET r.diagnosedDate = date(d.diagnosedDate),
            r.diagnosedBy = d.diagnosedBy,
            r.status = d.status,
            r.severity = d.severity,
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, diagnoses=diagnoses)
            return result.single()["created"]
    
    def load_prescriptions(self, prescriptions: List[Dict]) -> int:
        """Load prescription relationships."""
        query = """
        UNWIND $prescriptions AS p
        MATCH (patient:Patient {id: p.patientId})
        MATCH (med:Medication {id: p.medicationId})
        MERGE (patient)-[r:TAKES]->(med)
        SET r.dosage = p.dosage,
            r.frequency = p.frequency,
            r.prescribedDate = date(p.prescribedDate),
            r.prescribedBy = p.prescribedBy,
            r.durationDays = p.durationDays,
            r.refills = p.refills,
            r.status = p.status,
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, prescriptions=prescriptions)
            return result.single()["created"]
    
    def load_treatments(self, treatments: List[Dict]) -> int:
        """Load treatment relationships (Medication TREATS Disease)."""
        query = """
        UNWIND $treatments AS t
        MATCH (med:Medication {id: t.medicationId})
        MATCH (disease:Disease {id: t.diseaseId})
        MERGE (med)-[r:TREATS]->(disease)
        SET r.effectiveness = t.effectiveness,
            r.firstLine = t.firstLine,
            r.createdAt = datetime()
        RETURN count(*) as created
        """
        
        with self.driver.session() as session:
            result = session.run(query, treatments=treatments)
            return result.single()["created"]
    
    # =========================================================
    # BULK LOAD FROM FILE
    # =========================================================
    
    def load_from_json(self, filepath: str) -> Dict[str, int]:
        """
        Load all data from a JSON file.
        
        Expected JSON structure:
        {
            "patients": [...],
            "diseases": [...],
            "medications": [...],
            "diagnoses": [...],
            "prescriptions": [...],
            "treatments": [...],
            "documents": [...]
        }
        """
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        results = {}
        
        # Load entities
        if 'patients' in data:
            results['patients'] = self.load_patients(data['patients'])
        
        if 'diseases' in data:
            results['diseases'] = self.load_diseases(data['diseases'])
        
        if 'medications' in data:
            results['medications'] = self.load_medications(data['medications'])
        
        if 'documents' in data:
            results['documents'] = self.load_documents(data['documents'])
        
        # Load relationships
        if 'diagnoses' in data:
            results['diagnoses'] = self.load_diagnoses(data['diagnoses'])
        
        if 'prescriptions' in data:
            results['prescriptions'] = self.load_prescriptions(data['prescriptions'])
        
        if 'treatments' in data:
            results['treatments'] = self.load_treatments(data['treatments'])
        
        return results
    
    def setup_database(self, data_file: str = None):
        """
        Complete database setup:
        1. Clear existing data
        2. Create constraints
        3. Create indexes
        4. Load data
        """
        print("\n🏗️  Setting up Healthcare Knowledge Graph")
        print("=" * 50)
        
        # Clear
        print("\n1. Clearing database...")
        self.clear_database()
        
        # Constraints
        print("\n2. Creating constraints...")
        self.create_constraints()
        
        # Indexes
        print("\n3. Creating indexes...")
        self.create_indexes()
        
        # Load data
        if data_file and Path(data_file).exists():
            print(f"\n4. Loading data from {data_file}...")
            results = self.load_from_json(data_file)
            
            print("\n📊 Load Results:")
            for entity, count in results.items():
                print(f"   - {entity}: {count}")
        
        print("\n✅ Database setup complete!")
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get database statistics."""
        query = """
        MATCH (n)
        WITH labels(n)[0] as label, count(n) as count
        RETURN label, count
        ORDER BY count DESC
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            node_counts = {r["label"]: r["count"] for r in result}
        
        query = """
        MATCH ()-[r]->()
        WITH type(r) as type, count(r) as count
        RETURN type, count
        ORDER BY count DESC
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            rel_counts = {r["type"]: r["count"] for r in result}
        
        return {
            "nodes": node_counts,
            "relationships": rel_counts,
            "total_nodes": sum(node_counts.values()),
            "total_relationships": sum(rel_counts.values())
        }


# ============================================================
# SAMPLE DATA GENERATOR
# ============================================================

def generate_sample_data() -> Dict[str, List]:
    """Generate sample healthcare data for testing."""
    
    patients = [
        {"id": "P001", "name": "John Doe", "dateOfBirth": "1980-05-15",
         "gender": "Male", "bloodType": "O+", "allergies": ["Penicillin"], "active": True},
        {"id": "P002", "name": "Jane Smith", "dateOfBirth": "1992-08-22",
         "gender": "Female", "bloodType": "A-", "allergies": [], "active": True},
        {"id": "P003", "name": "Bob Wilson", "dateOfBirth": "1975-12-01",
         "gender": "Male", "bloodType": "B+", "allergies": ["Latex", "Sulfa"], "active": True},
        {"id": "P004", "name": "Alice Johnson", "dateOfBirth": "1988-03-10",
         "gender": "Female", "bloodType": "AB+", "allergies": ["Aspirin"], "active": True},
        {"id": "P005", "name": "Charlie Brown", "dateOfBirth": "1965-07-30",
         "gender": "Male", "bloodType": "O-", "allergies": [], "active": True},
    ]
    
    diseases = [
        {"id": "E11", "name": "Type 2 Diabetes Mellitus", "category": "Endocrine",
         "description": "Chronic metabolic disorder", "severity": "moderate"},
        {"id": "I10", "name": "Essential Hypertension", "category": "Cardiovascular",
         "description": "High blood pressure", "severity": "moderate"},
        {"id": "J45", "name": "Asthma", "category": "Respiratory",
         "description": "Chronic inflammatory disease", "severity": "mild"},
        {"id": "E78", "name": "Hyperlipidemia", "category": "Metabolic",
         "description": "Elevated blood lipids", "severity": "mild"},
    ]
    
    medications = [
        {"id": "6809", "name": "Metformin", "genericName": "Metformin Hydrochloride",
         "drugClass": "Biguanide", "dosageForms": ["Tablet", "Extended-release"]},
        {"id": "29046", "name": "Lisinopril", "genericName": "Lisinopril",
         "drugClass": "ACE Inhibitor", "dosageForms": ["Tablet"]},
        {"id": "36567", "name": "Atorvastatin", "genericName": "Atorvastatin Calcium",
         "drugClass": "Statin", "dosageForms": ["Tablet"]},
        {"id": "3992", "name": "Albuterol", "genericName": "Albuterol Sulfate",
         "drugClass": "Beta-2 Agonist", "dosageForms": ["Inhaler"]},
    ]
    
    diagnoses = [
        {"patientId": "P001", "diseaseId": "E11", "diagnosedDate": "2023-01-15",
         "diagnosedBy": "Dr. Sarah Johnson", "status": "active", "severity": "moderate"},
        {"patientId": "P001", "diseaseId": "I10", "diagnosedDate": "2023-03-20",
         "diagnosedBy": "Dr. Sarah Johnson", "status": "active", "severity": "mild"},
        {"patientId": "P002", "diseaseId": "J45", "diagnosedDate": "2020-06-10",
         "diagnosedBy": "Dr. Michael Chen", "status": "active", "severity": "moderate"},
        {"patientId": "P003", "diseaseId": "E11", "diagnosedDate": "2019-11-05",
         "diagnosedBy": "Dr. Emily White", "status": "active", "severity": "severe"},
        {"patientId": "P003", "diseaseId": "E78", "diagnosedDate": "2020-02-18",
         "diagnosedBy": "Dr. Emily White", "status": "active", "severity": "moderate"},
    ]
    
    prescriptions = [
        {"patientId": "P001", "medicationId": "6809", "dosage": "500mg",
         "frequency": "twice daily", "prescribedDate": "2023-01-15",
         "prescribedBy": "Dr. Sarah Johnson", "durationDays": 90, "refills": 3, "status": "active"},
        {"patientId": "P001", "medicationId": "29046", "dosage": "10mg",
         "frequency": "once daily", "prescribedDate": "2023-03-20",
         "prescribedBy": "Dr. Sarah Johnson", "durationDays": 90, "refills": 3, "status": "active"},
        {"patientId": "P002", "medicationId": "3992", "dosage": "2 puffs",
         "frequency": "as needed", "prescribedDate": "2020-06-10",
         "prescribedBy": "Dr. Michael Chen", "durationDays": 365, "refills": 12, "status": "active"},
        {"patientId": "P003", "medicationId": "6809", "dosage": "1000mg",
         "frequency": "twice daily", "prescribedDate": "2019-11-05",
         "prescribedBy": "Dr. Emily White", "durationDays": 90, "refills": 4, "status": "active"},
    ]
    
    treatments = [
        {"medicationId": "6809", "diseaseId": "E11", "effectiveness": "high", "firstLine": True},
        {"medicationId": "29046", "diseaseId": "I10", "effectiveness": "high", "firstLine": True},
        {"medicationId": "36567", "diseaseId": "E78", "effectiveness": "high", "firstLine": True},
        {"medicationId": "3992", "diseaseId": "J45", "effectiveness": "high", "firstLine": True},
    ]
    
    return {
        "patients": patients,
        "diseases": diseases,
        "medications": medications,
        "diagnoses": diagnoses,
        "prescriptions": prescriptions,
        "treatments": treatments
    }


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Healthcare Data Loader - Demonstration")
    print("=" * 60)
    
    try:
        loader = HealthcareDataLoader()
        
        # Option 1: Load from generated data
        print("\n📌 Loading Sample Data")
        print("-" * 40)
        
        data = generate_sample_data()
        
        # Setup constraints and indexes
        loader.create_constraints()
        loader.create_indexes()
        
        # Load entities
        print(f"\n   Patients loaded: {loader.load_patients(data['patients'])}")
        print(f"   Diseases loaded: {loader.load_diseases(data['diseases'])}")
        print(f"   Medications loaded: {loader.load_medications(data['medications'])}")
        
        # Load relationships
        print(f"   Diagnoses loaded: {loader.load_diagnoses(data['diagnoses'])}")
        print(f"   Prescriptions loaded: {loader.load_prescriptions(data['prescriptions'])}")
        print(f"   Treatments loaded: {loader.load_treatments(data['treatments'])}")
        
        # Statistics
        print("\n📊 Database Statistics")
        print("-" * 40)
        stats = loader.get_statistics()
        print(f"   Total nodes: {stats['total_nodes']}")
        print(f"   Total relationships: {stats['total_relationships']}")
        
        print("\n" + "=" * 60)
        print("✅ Data loader demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
